package solvers.linearsupport;

import globals.CentralStatics;
import globals.ExperimentTracker;
import globals.StopWatch;
import globals.Tuple;

import java.util.ArrayList;
import java.util.Iterator;

import problems.DiscreteActions;
import problems.DiscreteObservations;
import problems.Mopomdp;
import solvers.vi.BeliefSampler;
import solvers.vi.ValueIteratorPO;
//import value.CPruner;
import value.LazyScalarizedVector;
import value.SimpleVector;
import value.ValueAtBeliefVector;

public class ReusingPerseusSolver<S, A, O> implements SingleSolver<A> {

	public Mopomdp<S,A,O> problem;
	private int nSamples;
	private int maxIterations; 
	private double precision;
	public ArrayList<double[]> bSampleSet;
	private double[] initialBelief; 
	private ValueCollection<A> linsup;
	
	private int cnt;
	public static boolean LOG = false;
	
	public static ArrayList<SimpleVector<?>> valuelist; 
	public static ArrayList<Long> timelist; 
	
	public static final int TYPE_FULL = 42;
	public static final int TYPE_CPS  = 43;
	public static final int TYPE_CLOSE= 44;
	public static final int TYPE_NONE = 45;
	
	public int type;
	public static int asd;
	public static ArrayList<Double> listmax=new ArrayList<Double>();
	
	//resetMeasurement method creates two arraylist objects 
	// first ArrayList timelist is of Long type
	//second ArrayList valueList is of Simplevector type
	public static void resetMeasurement(boolean enforce){
		if(enforce){
			ReusingPerseusSolver.timelist = new ArrayList<Long>();
			ReusingPerseusSolver.valuelist= new ArrayList<SimpleVector<?>>();
		} else {
			if(ReusingPerseusSolver.timelist == null){
				ReusingPerseusSolver.timelist = new ArrayList<Long>();
			}
			if(ReusingPerseusSolver.valuelist == null){
				ReusingPerseusSolver.valuelist= new ArrayList<SimpleVector<?>>();
			}
		}
	}
	
	
	public ReusingPerseusSolver(Mopomdp<S,A,O> m, double[] iBelief, int nSamplez, int maxIter, double preciz, int typ){
		ReusingPerseusSolver.resetMeasurement(false);
		this.problem = m;
		this.bSampleSet = null;
		this.maxIterations = maxIter;
		this.nSamples = nSamplez;
		this.precision = preciz;
		this.initialBelief = iBelief;
		cnt = 0;
		if(typ==TYPE_CLOSE || typ == TYPE_FULL || typ==TYPE_NONE){
			type = typ;
		} else {
			type = TYPE_CPS;
		}
		linsup=null;
	}
	
	public ReusingPerseusSolver(Mopomdp<S,A,O> m, double[] iBelief, ArrayList<double[]>  bs, int maxIter, double preciz, int typ){
		this.problem = m;
		this.bSampleSet = bs;
		this.maxIterations = maxIter;
		this.nSamples = bs.size();
		this.precision = preciz;
		this.initialBelief = iBelief;
		cnt = 0;
		if(typ==TYPE_CLOSE || typ == TYPE_FULL){
			type = typ;
		} else {
			type = TYPE_CPS;
		}
		linsup=null;
	}
	
	public void initSamples(){
		if(bSampleSet == null){
			BeliefSampler<S,A,O> bs = new BeliefSampler<S,A,O>(this.problem);
			this.bSampleSet = bs.randomExplorationPerseus(nSamples);
			System.out.println("bSampleset initialized");
			
		}
	}
	
	public void setCPSreference(ValueCollection<A> ls){
		this.linsup = ls;
	}
	
	
	@SuppressWarnings("unchecked")
	public ArrayList<LazyScalarizedVector<A>> getFullLowerBound(double[] w){
		ArrayList<LazyScalarizedVector<A>> lb = new ArrayList<LazyScalarizedVector<A>>();
		if(this.linsup==null){
			System.err.println("CP list reference not set. Exiting...");
			System.exit(-42);
		}
		ArrayList<ValueVector<A>> vvList = linsup.getValueVectors();
		
		//System.out.println("VV size: "+vvList.size());
		Iterator<ValueVector<A>> cpit = vvList.iterator();
		while(cpit.hasNext()){
			ValueVector<A> vv = cpit.next();			
			ValueAtBeliefVector<A> vv2 = vv.p;
			//System.out.println(vv.p.retrieveTags());
			for(int j=0; j<vv2.alphas.size(); j++){
				//System.out.println(vv2.retrieveTags());
				////added code
				//System.out.println(vv2.retrieveTags().get(0));
				vv2.alphas.get(j).theTag=(A) vv2.retrieveTags().get(0);
				LazyScalarizedVector<A> lsv123=new LazyScalarizedVector<A>(vv2.alphas.get(j),w);
				
				
				
				//added code
				lb.add(lsv123);
				
			}
		}
		
		return lb;
	}
	
	public ArrayList<LazyScalarizedVector<A>> getClosestLowerBound(double[] w){
		ArrayList<LazyScalarizedVector<A>> lb = new ArrayList<LazyScalarizedVector<A>>();
		if(this.linsup==null){
			System.err.println("CP list reference not set. Exiting...");
			System.exit(-4132);
		}
		ArrayList<ValueVector<A>> vvList = linsup.getValueVectors();
		//System.out.println("VV size: "+vvList.size());
		Iterator<ValueVector<A>> cpit = vvList.iterator();
		double close = Double.POSITIVE_INFINITY;
		ValueVector<A> closest = null; 
		while(cpit.hasNext()){
			ValueVector<A> cp = cpit.next();
			double dist = CentralStatics.euclDistance(w, cp.getSearchWeights().get(0));
			if(dist<close){
				close = dist;
				closest = cp;
			}
		}
		if(closest == null){
			return lb;
		}
		ValueAtBeliefVector<A> vv = closest.p;
		for(int j=0; j<vv.alphas.size(); j++){
			lb.add(new LazyScalarizedVector<A>(vv.alphas.get(j), w));
		}
		return lb;
	}
	
	public ArrayList<LazyScalarizedVector<A>> getCPsLowerBound(CornerPoint<A> cp, double[] w){
		ArrayList<LazyScalarizedVector<A>> lb = new ArrayList<LazyScalarizedVector<A>>();
		for(int i=0; i<cp.valueVectors.size(); i++){
			ValueAtBeliefVector<A> vv = cp.valueVectors.get(i).p;
			for(int j=0; j<vv.alphas.size(); j++){
				lb.add(new LazyScalarizedVector<A>(vv.alphas.get(j), w));
			}
		}
		return lb;
	}
	
	public Tuple<Double, LazyScalarizedVector<A>> bestForB(ArrayList<LazyScalarizedVector<A>> list, double[] b){
		Tuple<Double, LazyScalarizedVector<A>> result = 
				new Tuple<Double, LazyScalarizedVector<A>>(Double.NEGATIVE_INFINITY, null);
		for(int i=0; i<list.size(); i++){
			LazyScalarizedVector<A> vec = list.get(i);
			double val = vec.linearScalValue(b);
			if(val>result.x.doubleValue()){
				result.x = val;
				result.y = vec;
			}
		}
		return result;
	}
	
	public ArrayList<LazyScalarizedVector<A>> bestForBeliefs(ArrayList<LazyScalarizedVector<A>> allList){
		ArrayList<LazyScalarizedVector<A>> lb = new ArrayList<LazyScalarizedVector<A>>(allList.size());
		
		for(int i=0; i<this.bSampleSet.size(); i++){
			double[] b = this.bSampleSet.get(i);
			Tuple<Double, LazyScalarizedVector<A>> auld  = this.bestForB(lb, b);
			Tuple<Double, LazyScalarizedVector<A>> novel = this.bestForB(allList, b);
			if(novel.x.doubleValue()>auld.x.doubleValue()){
				lb.add(novel.y);
			}
		}
		
		return lb;
	}
	
	@SuppressWarnings("unchecked")
	@Override
	public ValueAtBeliefVector<A> solveForW(CornerPoint<A> cp) {
		
		//System.out.println("yo");
		double[] w = cp.weightVector;
		
		this.initSamples();
		
		
		ArrayList<LazyScalarizedVector<A>> lb;
		if(type==TYPE_CPS){	
			lb = this.getCPsLowerBound(cp, w);
		} else if(type==TYPE_FULL){
			lb = this.getFullLowerBound(w);
		}  else if(type==TYPE_NONE){
			lb = new ArrayList<LazyScalarizedVector<A>>(1); 
			
		} else {
			lb = this.getClosestLowerBound(w);
		}
		
		//System.out.println(">>>> LB-size: "+lb.size());
		LazyScalarizedVector<A> lsv = new LazyScalarizedVector<A>(this.problem.misereMatrix(), w);
		lb.add(lsv);
		
		lb = this.bestForBeliefs(lb);
		//System.out.println(">>>> LB-size bis: "+lb.size());
		
		

		ValueIteratorPO<LazyScalarizedVector<A>,S,A,O> vi = 
				new ValueIteratorPO<LazyScalarizedVector<A>,S,A,O>(this.problem, lb);
		
		
		try{
			DiscreteActions<A,S> dacts   = (DiscreteActions<A,S>) problem.getActionSet();
			DiscreteObservations<O> dobs = (DiscreteObservations<O>) problem.getObservationSet();
			vi.perseus(bSampleSet, dacts, dobs, problem.getStateSet(), this.precision, this.maxIterations);
			ValueAtBeliefVector<A> vb = new ValueAtBeliefVector<A>(w.length, problem.currentBelief(), vi.alphas);
			
			
			
			
			
			//System.out.println("~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~");
			/////********System.out.println("Checking value at belief vector valuevec"+vb.getValue()[0]+" "+vb.getValue()[1]);
			Double valvb=vb.linearScalValue(w);
			//System.out.println("Scalarized valvb:   "+valvb);
			
			//System.out.println("~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~");
			//System.out.println("~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~");
			
			//this.listmax.add(t.x);
			this.listmax.add(valvb);

			
			/////////////////////////////////////////
			
			
			
			if(LOG){
				StopWatch sw = StopWatch.getGlobalStopWatch();
				sw.stop();
				
				SimpleVector<A> s = new SimpleVector<A>(vb);
				ReusingPerseusSolver.valuelist.add(s);
				ReusingPerseusSolver.timelist.add(sw.totalElapsedTime());
				sw.start();
			}
			return vb;
		} catch(Exception e) {
			e.printStackTrace();
			System.out.println(e);
			System.err.println("Trying to call discrete POMDP method for non-discrete problem.");
			System.err.println("Quitting...");
			System.exit(77);
			return null;
		}
	}

	@SuppressWarnings("unchecked")
	@Override
	public ValueAtBeliefVector<A> solveForW(double[] w) {
		
		this.initSamples();
		
		ArrayList<LazyScalarizedVector<A>> lb;

		lb = new ArrayList<LazyScalarizedVector<A>>();
		if(type==TYPE_CPS){	
			lb = this.getFullLowerBound(w);
		} else if(type==TYPE_FULL){
			lb = this.getFullLowerBound(w);
		} else if(type==TYPE_NONE){
			//do nothing
			//lb = new ArrayList<LazyScalarizedVector<A>>(1); 
			//LazyScalarizedVector<A> lsv = new LazyScalarizedVector<A>(this.problem.misereMatrix(), w);
			//lb.add(lsv);
		} else {
			lb = this.getClosestLowerBound(w);
		}
		
		//System.out.println("LazyScalarized lb size:  "+lb.size());
		
		LazyScalarizedVector<A> lsv = new LazyScalarizedVector<A>(this.problem.misereMatrix(), w);
		lb.add(lsv);
		//System.out.println("lsv weights:  "+w[0]+"  "+w[1]+" "+w[2]);
		//System.out.println("LazyScalarized lsv tag:  "+lsv.theTag);
		
		lb = this.bestForBeliefs(lb);
		
		
		
		ValueIteratorPO<LazyScalarizedVector<A>,S,A,O> vi = 
				new ValueIteratorPO<LazyScalarizedVector<A>,S,A,O>(this.problem, lb);
		try{
			DiscreteActions<A,S> dacts   = (DiscreteActions<A,S>) problem.getActionSet();
			DiscreteObservations<O> dobs = (DiscreteObservations<O>) problem.getObservationSet();
			vi.perseus(bSampleSet, dacts, dobs, problem.getStateSet(), this.precision, this.maxIterations);
			ValueAtBeliefVector<A> vb = new ValueAtBeliefVector<A>(w.length, this.initialBelief, vi.alphas);
			
			//System.out.println(" vb tag:"+vb.retrieveTags());
			
			if(LOG){
				StopWatch sw = StopWatch.getGlobalStopWatch();
				sw.stop();
				
				SimpleVector<A> s = new SimpleVector<A>(vb);
				ReusingPerseusSolver.valuelist.add(s);
				ReusingPerseusSolver.timelist.add(sw.totalElapsedTime());
				sw.start();
				
				
			}
			return vb;
		} catch(Exception e) {
			e.printStackTrace();
			System.err.println("Trying to call discrete POMDP method for non-discrete problem.");
			System.err.println("Quitting...");
			System.exit(77);
			return null;
		}
		
	}
}
